import unittest
import warnings

import numpy as np
import pandas as pd
from joblib.externals.loky import get_reusable_executor
from skbase.utils.dependencies import _check_soft_dependencies

from pgmpy import config
from pgmpy.estimators import ExpectationMaximization as EM
from pgmpy.factors.discrete import TabularCPD
from pgmpy.models import DiscreteBayesianNetwork
from pgmpy.utils import compat_fns, get_example_model


class TestEM(unittest.TestCase):
    def setUp(self):
        self.model1 = get_example_model("cancer")
        self.data1 = self.model1.simulate(int(1e4), seed=42)

        self.model2 = DiscreteBayesianNetwork(self.model1.edges(), latents={"Smoker"})
        self.model2.add_cpds(*self.model1.cpds)
        self.data2 = self.model2.simulate(int(1e4), seed=42)

    def test_get_parameters(self):
        # All observed
        est = EM(self.model1, self.data1)
        cpds = est.get_parameters(seed=42, n_jobs=1, show_progress=False)
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

        # Latent variables
        est = EM(self.model2, self.data2)
        cpds = est.get_parameters(seed=42, n_jobs=1, show_progress=False)
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model2.get_cpds(var)

            if "Smoker" in orig_cpd.variables:
                orig_cpd.state_names["Smoker"] = [1, 0]
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

    def test_get_parameters_smoothing_k2(self):
        # All observed
        est = EM(self.model1, self.data1)
        cpds = est.get_parameters(
            seed=42,
            n_jobs=1,
            apply_smoothing=True,
            prior_type="k2",
            show_progress=False,
        )
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

        # Latent variables
        est = EM(self.model2, self.data2)
        cpds = est.get_parameters(
            seed=42,
            n_jobs=1,
            apply_smoothing=True,
            prior_type="k2",
            show_progress=False,
        )
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model2.get_cpds(var)

            if "Smoker" in orig_cpd.variables:
                orig_cpd.state_names["Smoker"] = [1, 0]
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

    def test_get_parameters_smoothing_bdeu(self):
        # All observed
        est = EM(self.model1, self.data1)
        cpds = est.get_parameters(
            seed=42,
            n_jobs=1,
            apply_smoothing=True,
            prior_type="bdeu",
            equivalent_sample_size=1,
            show_progress=False,
        )
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

        # Latent variables
        est = EM(self.model2, self.data2)
        cpds = est.get_parameters(
            seed=42,
            n_jobs=1,
            apply_smoothing=True,
            prior_type="bdeu",
            equivalent_sample_size=1,
            show_progress=False,
        )
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model2.get_cpds(var)

            if "Smoker" in orig_cpd.variables:
                orig_cpd.state_names["Smoker"] = [1, 0]
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

    def test_get_parameters_initial_cpds(self):
        # All observed. Specify initial CPDs.
        est = EM(self.model1, self.data1)
        smoker_initial = TabularCPD(
            "Smoker", 2, [[0.1], [0.9]], state_names={"Smoker": ["True", "False"]}
        )
        cpds = est.get_parameters(
            init_cpds={"Smoker": smoker_initial}, seed=42, n_jobs=1, show_progress=False
        )
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

        # With latents. Specify initial CPDs only for latent.
        est = EM(self.model2, self.data2)
        cpds = est.get_parameters(
            init_cpds={"Smoker": smoker_initial}, seed=42, n_jobs=1, show_progress=False
        )
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            if "Smoker" in orig_cpd.variables:
                orig_cpd.state_names["Smoker"] = [1, 0]

            # The latent variable doesn't converge to the true value when
            # the initial CPD is specified.
            if orig_cpd.variables[0] == "Smoker":
                self.assertTrue(
                    np.allclose(est_cpd.values, np.array([0.123, 0.877]), atol=0.01)
                )
            else:
                self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

        # With latents. Specify initial CPDs for both latents and observed.
        est = EM(self.model2, self.data2)
        xray_initial = TabularCPD(
            variable="Xray",
            variable_card=2,
            values=[[0.1, 0.8], [0.9, 0.2]],
            evidence=["Cancer"],
            evidence_card=[2],
            state_names={"Xray": ["positive", "negative"], "Cancer": ["True", "False"]},
        )
        cpds = est.get_parameters(
            init_cpds={"Smoker": smoker_initial, "Xray": xray_initial},
            seed=42,
            n_jobs=1,
            show_progress=False,
        )

        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            if "Smoker" in orig_cpd.variables:
                orig_cpd.state_names["Smoker"] = [1, 0]

            # The latent variable doesn't converge to the true value when
            # the initial CPD is specified.
            if orig_cpd.variables[0] == "Smoker":
                self.assertTrue(
                    np.allclose(est_cpd.values, np.array([0.123, 0.877]), atol=0.01)
                )
            elif orig_cpd.variables[0] == "Xray":
                self.assertTrue(
                    np.allclose(
                        est_cpd.values,
                        np.array([[0.799, 0.093], [0.201, 0.907]]),
                        atol=0.01,
                    )
                )
            else:
                self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

    def test_em_init_missing_data_handling(self):
        df = pd.DataFrame(
            {"A": [1, 2, 3], "B": [None, None, None], "C": [1, None, 3], "D": [4, 5, 6]}
        )

        with warnings.catch_warnings(record=True):
            warnings.simplefilter("always")
            est = EM(self.model1, df)

        # Data shape and column removal
        self.assertEqual(est.data.shape, (2, 3))
        self.assertNotIn("B", est.data.columns)

    def test_get_parameters_random_init_cpds(self):
        est = EM(self.model1, self.data1)
        cpds = est.get_parameters(
            init_cpds="random", seed=42, n_jobs=1, show_progress=False
        )
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

    def test_get_parameters_uniform_init_cpds(self):
        est = EM(self.model1, self.data1)
        cpds = est.get_parameters(init_cpds="uniform", n_jobs=1, show_progress=False)
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

    def test_get_parameters_node_specific_ess_bdeu(self):
        """Test EM with node-specific equivalent_sample_size dict for BDeu."""
        # All observed
        est = EM(self.model1, self.data1)
        ess_dict = {"Smoker": 10, "Cancer": 5, "Xray": 8}
        cpds = est.get_parameters(
            seed=42,
            n_jobs=1,
            apply_smoothing=True,
            prior_type="bdeu",
            equivalent_sample_size=ess_dict,
            show_progress=False,
        )
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

        # With latent variables
        est = EM(self.model2, self.data2)
        cpds = est.get_parameters(
            seed=42,
            n_jobs=1,
            apply_smoothing=True,
            prior_type="bdeu",
            equivalent_sample_size=ess_dict,
            show_progress=False,
        )
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model2.get_cpds(var)

            if "Smoker" in orig_cpd.variables:
                orig_cpd.state_names["Smoker"] = [1, 0]
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

    def test_get_parameters_ess_dict_vs_scalar(self):
        """Test that uniform ESS dict matches scalar ESS."""
        ess_value = 7
        ess_dict = {"Smoker": ess_value, "Cancer": ess_value, "Xray": ess_value}

        est_scalar = EM(self.model1, self.data1)
        cpds_scalar = est_scalar.get_parameters(
            seed=42,
            n_jobs=1,
            apply_smoothing=True,
            prior_type="bdeu",
            equivalent_sample_size=ess_value,
            show_progress=False,
        )

        est_dict = EM(self.model1, self.data1)
        cpds_dict = est_dict.get_parameters(
            seed=42,
            n_jobs=1,
            apply_smoothing=True,
            prior_type="bdeu",
            equivalent_sample_size=ess_dict,
            show_progress=False,
        )

        # Results should be identical
        for cpd_scalar, cpd_dict in zip(
            sorted(cpds_scalar, key=lambda x: x.variables[0]),
            sorted(cpds_dict, key=lambda x: x.variables[0]),
        ):
            self.assertTrue(cpd_scalar.__eq__(cpd_dict, atol=1e-6))

    def tearDown(self):
        del self.model1
        del self.model2
        del self.data1
        del self.data2

        get_reusable_executor().shutdown(wait=True)


@unittest.skipUnless(
    _check_soft_dependencies("torch", severity="none"),
    reason="execute only if required dependency present",
)
class TestEMTorch(TestEM):
    def setUp(self):
        config.set_backend("torch")

        self.model1 = get_example_model("cancer")
        self.data1 = self.model1.simulate(int(1e4), seed=42)

        self.model2 = DiscreteBayesianNetwork(self.model1.edges(), latents={"Smoker"})
        self.model2.add_cpds(*self.model1.cpds)
        self.data2 = self.model2.simulate(int(1e4), seed=42)

    def test_get_parameters(self):
        est = EM(self.model1, self.data1)
        cpds = est.get_parameters(seed=42, n_jobs=1, show_progress=False)
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

        est = EM(self.model2, self.data2)
        cpds = est.get_parameters(seed=42, n_jobs=1, show_progress=False)
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model2.get_cpds(var)

            if "Smoker" in orig_cpd.variables:
                orig_cpd.state_names["Smoker"] = [1, 0]

            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

    def test_get_parameters_initial_cpds(self):
        # All observed. Specify initial CPDs.
        est = EM(self.model1, self.data1)
        smoker_initial = TabularCPD(
            "Smoker", 2, [[0.1], [0.9]], state_names={"Smoker": ["True", "False"]}
        )
        cpds = est.get_parameters(
            init_cpds={"Smoker": smoker_initial}, seed=42, n_jobs=1, show_progress=False
        )
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

        # With latents. Specify initial CPDs only for latent.
        est = EM(self.model2, self.data2)
        cpds = est.get_parameters(
            init_cpds={"Smoker": smoker_initial}, seed=42, n_jobs=1, show_progress=False
        )
        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            if "Smoker" in orig_cpd.variables:
                orig_cpd.state_names["Smoker"] = [1, 0]

            # The latent variable doesn't converge to the true value when
            # the initial CPD is specified.
            if orig_cpd.variables[0] == "Smoker":
                self.assertTrue(
                    np.allclose(
                        compat_fns.to_numpy(est_cpd.values),
                        np.array([0.123, 0.877]),
                        atol=0.01,
                    )
                )
            else:
                self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

        # With latents. Specify initial CPDs for both latents and observed.
        est = EM(self.model2, self.data2)
        xray_initial = TabularCPD(
            variable="Xray",
            variable_card=2,
            values=[[0.1, 0.8], [0.9, 0.2]],
            evidence=["Cancer"],
            evidence_card=[2],
            state_names={"Xray": ["positive", "negative"], "Cancer": ["True", "False"]},
        )
        cpds = est.get_parameters(
            init_cpds={"Smoker": smoker_initial, "Xray": xray_initial},
            seed=42,
            n_jobs=1,
            show_progress=False,
        )

        for est_cpd in cpds:
            var = est_cpd.variables[0]
            orig_cpd = self.model1.get_cpds(var)
            if "Smoker" in orig_cpd.variables:
                orig_cpd.state_names["Smoker"] = [1, 0]

            # The latent variable doesn't converge to the true value when
            # the initial CPD is specified.
            if orig_cpd.variables[0] == "Smoker":
                self.assertTrue(
                    np.allclose(
                        compat_fns.to_numpy(est_cpd.values),
                        np.array([0.123, 0.877]),
                        atol=0.01,
                    )
                )
            elif orig_cpd.variables[0] == "Xray":
                self.assertTrue(
                    np.allclose(
                        compat_fns.to_numpy(est_cpd.values),
                        np.array([[0.799, 0.093], [0.201, 0.907]]),
                        atol=0.01,
                    )
                )
            else:
                self.assertTrue(orig_cpd.__eq__(est_cpd, atol=0.1))

    def tearDown(self):
        del self.model1
        del self.model2
        del self.data1
        del self.data2

        get_reusable_executor().shutdown(wait=True)

        config.set_backend("numpy")
